import torch
import torch.nn.functional as F

"""
不同解码策略（贪婪搜索、随机采样、核采样）的简单实现，展示它们如何体现"接受率"的概念

1. 贪婪搜索 (Greedy Search)
接受率特点: 对最高概率token为100%，其他为0%

行为: 每次选择概率最高的token，结果确定且可预测

适用场景: 需要确定性输出的场景

2. 随机采样 (Random Sampling)
接受率特点: 每个token的概率值即为其接受率

行为: 完全按照概率分布随机选择，多样性最高

适用场景: 需要高度创造性输出的场景

3. 核采样 (Top-p Sampling)
接受率特点: 对累积概率超过p的最小候选集中的token重新分配接受率

行为: 只在高质量候选集中随机选择，平衡质量与多样性

适用场景: 大多数文本生成任务的最佳选择

4. Top-k采样 (Top-k Sampling)
接受率特点: 对概率排名前k的token重新分配接受率

行为: 只在前k个候选token中随机选择

适用场景: 需要控制候选集大小的场景

接受率概念总结
在这三种策略中，"接受率"体现在：

贪婪搜索: 接受率是二元的 - 最高概率token为100%，其他为0%

随机采样: 接受率是连续的 - 每个token的概率值就是它的接受率

核采样/Top-k采样: 接受率是条件性的 - 只对筛选后的候选集重新分配接受率

这些策略提供了不同的权衡：

确定性 vs 随机性: 贪婪搜索完全确定，随机采样完全随机

质量 vs 多样性: 贪婪搜索质量最高但多样性最低，随机采样多样性最高但可能包含低质量输出

可控性: 核采样和Top-k采样通过参数(top_p, top_k)提供可控的随机性

在实际应用中，核采样(top-p)通常是最佳选择，因为它能动态地选择高质量的候选集，同时在候选集内保持随机性，平衡了输出质量和多样性。
"""
def greedy_search(logits):
    """
    贪婪搜索策略
    接受率: 对最高概率token为100%，其他为0%
    """
    # 直接选择概率最高的token
    next_token = torch.argmax(logits, dim=-1)
    return next_token


def random_sampling(logits, temperature=1.0):
    """
    随机采样策略
    接受率: 每个token的概率值即为其接受率
    """
    # 应用温度参数
    scaled_logits = logits / temperature
    # 转换为概率分布
    probs = F.softmax(scaled_logits, dim=-1)
    # 从分布中随机采样
    next_token = torch.multinomial(probs, num_samples=1)
    return next_token


def top_p_sampling(logits, top_p=0.9, temperature=1.0):
    """
    核采样 (Top-p采样) 策略
    接受率: 对累积概率超过p的最小候选集中的token重新分配接受率
    """
    # 应用温度参数
    scaled_logits = logits / temperature
    probs = F.softmax(scaled_logits, dim=-1)

    # 对概率进行排序
    sorted_probs, sorted_indices = torch.sort(probs, descending=True)

    # 计算累积概率
    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

    # 移除累积概率超过top_p的token
    sorted_indices_to_remove = cumulative_probs > top_p
    # 确保至少保留一个token
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = False

    # 将需要移除的token概率置为0
    indices_to_remove = sorted_indices_to_remove.scatter(
        -1, sorted_indices, sorted_indices_to_remove
    )
    probs = probs.masked_fill(indices_to_remove, 0.0)

    # 重新归一化概率分布
    if torch.sum(probs) > 0:
        probs = probs / torch.sum(probs)
    else:
        # 如果所有概率都为0，回退到均匀分布
        probs = torch.ones_like(probs) / probs.size(-1)

    # 从新的分布中采样
    next_token = torch.multinomial(probs, num_samples=1)
    return next_token


def top_k_sampling(logits, top_k=50, temperature=1.0):
    """
    Top-k采样策略
    接受率: 对概率排名前k的token重新分配接受率
    """
    # 应用温度参数
    scaled_logits = logits / temperature
    probs = F.softmax(scaled_logits, dim=-1)

    # 获取top-k概率值和索引
    top_k_probs, top_k_indices = torch.topk(probs, top_k, dim=-1)

    # 创建一个新的分布，将不在top-k内的概率置为0
    probs = torch.zeros_like(probs)
    probs.scatter_(-1, top_k_indices, top_k_probs)

    # 重新归一化概率分布
    probs = probs / torch.sum(probs, dim=-1, keepdim=True)

    # 从新的分布中采样
    next_token = torch.multinomial(probs, num_samples=1)
    return next_token


# 演示函数
def demonstrate_sampling_strategies():
    """演示不同采样策略的接受率"""
    # 模拟模型输出的logits (假设词汇表大小为10)
    logits = torch.tensor([[2.0, 1.0, 0.5, 0.2, 0.1, -0.5, -1.0, -1.5, -2.0, -2.5]])
    probs = F.softmax(logits, dim=-1)

    print("原始概率分布:")
    for i, prob in enumerate(probs[0]):
        print(f"Token {i}: {prob:.4f}")

    print("\n=== 不同采样策略演示 ===\n")

    # 贪婪搜索
    greedy_token = greedy_search(logits)
    print(f"贪婪搜索选择的Token: {greedy_token.item()}")
    print(f"接受率: Token {greedy_token.item()} = 100%, 其他 = 0%")

    # 随机采样
    random_token = random_sampling(logits, temperature=1.0)
    print(f"\n随机采样选择的Token: {random_token.item()}")
    print("接受率: 每个token的概率值即为其接受率")

    # Top-p采样
    top_p_token = top_p_sampling(logits, top_p=0.9, temperature=1.0)
    print(f"\nTop-p采样选择的Token: {top_p_token.item()}")

    # 计算Top-p采样后的概率分布
    scaled_logits = logits / 1.0
    probs = F.softmax(scaled_logits, dim=-1)
    sorted_probs, sorted_indices = torch.sort(probs, descending=True)
    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
    sorted_indices_to_remove = cumulative_probs > 0.9
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = False

    indices_to_remove = sorted_indices_to_remove.scatter(
        -1, sorted_indices, sorted_indices_to_remove
    )
    filtered_probs = probs.masked_fill(indices_to_remove, 0.0)
    if torch.sum(filtered_probs) > 0:
        filtered_probs = filtered_probs / torch.sum(filtered_probs)

    print("Top-p采样后的接受率分布:")
    for i, prob in enumerate(filtered_probs[0]):
        if prob > 0:
            print(f"Token {i}: {prob:.4f} (接受)")
        else:
            print(f"Token {i}: {prob:.4f} (拒绝)")

    # Top-k采样
    top_k_token = top_k_sampling(logits, top_k=3, temperature=1.0)
    print(f"\nTop-k采样选择的Token: {top_k_token.item()}")

    # 计算Top-k采样后的概率分布
    probs = F.softmax(logits, dim=-1)
    top_k_probs, top_k_indices = torch.topk(probs, 3, dim=-1)
    filtered_probs = torch.zeros_like(probs)
    filtered_probs.scatter_(-1, top_k_indices, top_k_probs)
    filtered_probs = filtered_probs / torch.sum(filtered_probs, dim=-1, keepdim=True)

    print("Top-k采样后的接受率分布:")
    for i, prob in enumerate(filtered_probs[0]):
        if prob > 0:
            print(f"Token {i}: {prob:.4f} (接受)")
        else:
            print(f"Token {i}: {prob:.4f} (拒绝)")


# 多次运行演示，观察不同策略的行为
def run_multiple_times():
    """多次运行不同策略，观察其行为差异"""
    logits = torch.tensor([[3.0, 2.0, 1.0, 0.0, -1.0, -2.0, -3.0]])
    probs = F.softmax(logits, dim=-1)

    print("原始概率分布:")
    for i, prob in enumerate(probs[0]):
        print(f"Token {i}: {prob:.4f}")

    print("\n=== 多次运行结果 ===\n")

    # 贪婪搜索 (每次结果相同)
    greedy_results = []
    for _ in range(10):
        greedy_results.append(greedy_search(logits).item())
    print(f"贪婪搜索结果: {greedy_results}")

    # 随机采样 (每次结果可能不同)
    random_results = []
    for _ in range(10):
        random_results.append(random_sampling(logits).item())
    print(f"随机采样结果: {random_results}")

    # Top-p采样 (每次结果可能不同，但只在高质量候选集中)
    top_p_results = []
    for _ in range(10):
        top_p_results.append(top_p_sampling(logits, top_p=0.9).item())
    print(f"Top-p采样结果: {top_p_results}")

    # Top-k采样 (每次结果可能不同，但只在前k个候选集中)
    top_k_results = []
    for _ in range(10):
        top_k_results.append(top_k_sampling(logits, top_k=3).item())
    print(f"Top-k采样结果: {top_k_results}")


if __name__ == "__main__":
    print("=" * 50)
    print("接受率三种策略演示")
    print("=" * 50)

    demonstrate_sampling_strategies()

    print("\n" + "=" * 50)
    print("多次运行对比")
    print("=" * 50)

    run_multiple_times()